#!/bin/bash

export OUTPUT_DIR='/app/data/output/real_estate_construction-0__1B'
export SWANLAB_API_KEY=GjTPmINWI217Tmfsxh98x
export SWANLAB_LOG_DIR="${OUTPUT_DIR}/logs"
mkdir -p "$SWANLAB_LOG_DIR"
if [ -z "$MODELSCOPE_CACHE" ]; then
    export MODELSCOPE_CACHE="$HOME/.cache/modelscope/hub"
fi
export current_time=$(date +"%Y%m%d%H%M%S")

CUDA_VISIBLE_DEVICES=0,1 \
NPROC_PER_NODE=2 \
OMP_NUM_THREADS=4 \
swift pt \
    --model /root/.cache/pminfo/AiLab-0.1B \
    --train_type full \
    --cached_dataset "${OUTPUT_DIR}/pretrain_cached_dataset" \
    --output_dir "${OUTPUT_DIR}" \
    --add_version false  \
    --eval_strategy steps \
    --eval_steps 1000 \
    --save_steps 1000 \
    --save_total_limit 3 \
    --num_train_epochs 10 \
    --max_length 8192 \
    --per_device_train_batch_size 4 \
    --per_device_eval_batch_size 4 \
    --gradient_accumulation_steps 4 \
    --learning_rate 1e-5 \
    --warmup_ratio 0.05 \
    --dataloader_num_workers 8 \
    --dataset_num_proc 8 \
    --report_to swanlab \
    --swanlab_project AiLab \
    --swanlab_exp_name "estate_${current_time}" \
    --attn_impl flash_attn \
    --dataloader_drop_last true \
    --deepspeed zero2
